import pandas as pd
import re

prefix = '../data/'
columns = ['label-coarse', 'label-fine', 'text']

class TREC6:
    def load_data(self,filepath):
        data = []
        i = 0
        with open(filepath) as f:
            data_str = f.read()
            if data_str:
                for line in data_str.split('\n'):
                    if line:
                        # label_coarse = line.split(':')[0]
                        label_coarse = ''.join(re.findall(r"^\b[A-Z]*\b", line))
                        label_fine = ''.join(re.findall(r"\b(?<=:)[a-z]*\b", line))
                        sentence = ''.join(re.findall(r"(?<=:).*[A-Z]*$", line))
                        text = ' '.join(sentence.split(' ')[1:])
                        data.append([label_coarse, label_fine, text])
        return data

    def train_data(self):
        filepath = prefix + 'trec/train.txt'
        data = self.load_data(filepath)
        train_data = pd.DataFrame(data, columns=columns)
        return train_data

    def test_data(self):
        filepath = prefix + 'trec/test.txt'
        data = self.load_data(filepath)
        test_data = pd.DataFrame(data, columns=columns)
        return test_data

    def label_encoder(self, data):
        """
        将标签数值化
        :param data:
        :return:
        """
        _COARSE_LABELS = {"DESC": 0, "ENTY": 1, "ABBR": 2, "HUM": 3, "NUM": 4, "LOC": 5}
        _FINE_LABELS = {
            "manner":0,
            "cremat":1,
            "animal":2,
            "exp":3,
            "ind":4,
            "gr":5,
            "title":6,
            "def":7,
            "date":8,
            "reason":9,
            "event":10,
            "state":11,
            "desc":12,
            "count":13,
            "other":14,
            "letter":15,
            "religion":16,
            "food":17,
            "country":18,
            "color":19,
            "termeq":20,
            "city":21,
            "body":22,
            "dismed":23,
            "mount":24,
            "money":25,
            "product":26,
            "period":27,
            "substance":28,
            "sport":29,
            "plant":30,
            "techmeth":31,
            "volsize":32,
            "instru":33,
            "abb":34,
            "speed":35,
            "word":36,
            "lang":37,
            "perc":38,
            "code":39,
            "dist":40,
            "temp":41,
            "symbol":42,
            "ord":43,
            "veh":44,
            "weight":45,
            "currency":46
        }


        data['label-coarse'] = data['label-coarse'].map(lambda x: _COARSE_LABELS[x])
        data['label-fine'] = data['label-fine'].map(lambda x: _FINE_LABELS[x])
        return data

    def TREC6_Data(self):
        train_data = self.train_data()
        test_data = self.test_data()
        train_data = self.label_encoder(train_data)
        test_data = self.label_encoder(test_data)
        return train_data, test_data

    def bert_load_data(self):
        '''
        将数据转换为bert4keras的Datagenerator输入的形式
        :return: [(label, text)]
        '''
        train_data, test_data = self.TREC6_Data()
        train = []
        test = []
        for i in range(len(train_data)):
            train.append((train_data.iloc[i][0], train_data.iloc[i][-1]))
        for i in range(len(test_data)):
            test.append((test_data.iloc[i][0], test_data.iloc[i][-1]))
        return train, test


class MedQA:
    def load_data(self):
        train_data = pd.read_csv(prefix + 'medQA/medQA.train.csv')
        test_data = pd.read_csv(prefix + 'medQA/medQA.test.csv')
        val_data = pd.read_csv(prefix + 'medQA/medQA.valid.csv')
        train_data = self.findchinese(train_data)
        test_data = self.findchinese(test_data)
        val_data = self.findchinese(val_data)
        return train_data, test_data, val_data

    def train_data(self):
        train_data = pd.read_csv(prefix + 'medQA/medQA.train.csv')
        train_data = self.findchinese(train_data)
        return train_data

    def findchinese(self, data):
        '''
        使用正则表达式提取每一句子中的中文，去除其他的字符
        :return:
        '''
        text = []
        for i in data['question']:
            pattern = re.compile(r'[^\u4e00-\u9fa5]')
            sentence = re.sub(pattern, '', i)
            text.append(sentence)
        data['question'] = text
        return data

    def label_encoder(self, data):
        target_names = {'内科':0,
                        '外科':1,
                        '妇产科':2,
                        '儿科':3,
                        '皮肤性病科':4,
                        '五官科':5,
                        '肿瘤科':6,
                        '心理健康科':7,
                        '中医科':8,
                        '传染科':9,
                        '整形美容科':10,
                        '美容':11,
                        '药品':12,
                        '辅助检查科':13,
                        '保健养生':14,
                        '康复医学科':15,
                        '家居环境':16,
                        '子女教育':17,
                        '营养保健科':18,
                        '运动瘦身':19,
                        '遗传':20,
                        '体检科':21,
                        '其他科室':22}
        data['label'] = data['label'].map(lambda x: target_names[x])
        return data

    def MedQAData(self):
        traindata, testdata, valdata = self.load_data()
        train_data = self.label_encoder(traindata)
        test_data = self.label_encoder(testdata)
        val_data = self.label_encoder(valdata)
        return train_data, test_data, val_data

    def medqa_load_data(self):
        '''
        将数据转换为bert4keras的datagenerator输入的形式
        :return: [(label, text)]
        '''
        train_data, test_data, val_data = self.MedQAData()
        train = []
        val = []
        test = []
        for i in range(len(train_data)):
            train.append((train_data.iloc[i][1], train_data.iloc[i][0]))
        for i in range(len(val_data)):
            val.append((val_data.iloc[i][1], val_data.iloc[i][0]))
        for i in range(len(test_data)):
            test.append((test_data.iloc[i][1], test_data.iloc[i][0]))
        return train, test, val
